package xyz.lwm.lazycat.utility;

import java.util.HashMap;
import java.util.function.Function;

public class CustomKeyMap<K, V> extends HashMap<K, V> {

    private final Function<K, K> keyMapper;

    public CustomKeyMap(Function<K, K> keyMapper) {
        this.keyMapper = keyMapper;
    }

    @Override
    @SuppressWarnings("unchecked")
    public V get(Object key) {
        return super.get(keyMapper.apply((K) key));
    }

    @Override
    @SuppressWarnings("unchecked")
    public boolean containsKey(Object key) {
        return super.containsKey(keyMapper.apply( (K) key ));
    }

    @Override
    public V put(K key, V value) {
        return super.put(keyMapper.apply(key), value);
    }

    @Override
    @SuppressWarnings("unchecked")
    public V remove(Object key) {
        return super.remove(keyMapper.apply((K) key));
    }

}
